# 小球碰撞其他小球
import pygame
import sys
import random
import math

# 初始化Pygame
pygame.init()

# 设置窗口大小
screen_width = 800
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))

# 设置标题
pygame.display.set_caption("多彩小球游戏")

# 小球的位置、速度和半径
balls = []
num_balls = 5
ball_radius_min = 10
ball_radius_max = 20
speed_range = [-2, -1, 1, 2]


def calculate_new_velocities(ball_a, ball_b):
    # 计算两球中心之间的向量
    dx = ball_b['pos'][0] - ball_a['pos'][0]
    dy = ball_b['pos'][1] - ball_a['pos'][1]
    distance = math.sqrt(dx ** 2 + dy ** 2)

    # 计算碰撞点的法线方向
    normal = (dx / distance, dy / distance)
    tangent = (-normal[1], normal[0])

    # 计算沿法线方向的速度分量
    va_n = dot(ball_a['speed'], normal)
    vb_n = dot(ball_b['speed'], normal)
    va_t = dot(ball_a['speed'], tangent)
    vb_t = dot(ball_b['speed'], tangent)

    # 根据弹性碰撞公式计算新的速度
    va_n_new = ((va_n * (ball_a['radius'] - ball_b['radius']) + 2 * ball_b['radius'] * vb_n) /
                (ball_a['radius'] + ball_b['radius']))
    vb_n_new = ((vb_n * (ball_b['radius'] - ball_a['radius']) + 2 * ball_a['radius'] * va_n) /
                (ball_a['radius'] + ball_b['radius']))

    # 合成新的速度向量
    ball_a['speed'] = [va_t * tangent[0] + va_n_new * normal[0],
                       va_t * tangent[1] + va_n_new * normal[1]]
    ball_b['speed'] = [vb_t * tangent[0] + vb_n_new * normal[0],
                       vb_t * tangent[1] + vb_n_new * normal[1]]


def dot(v1, v2):
    return v1[0] * v2[0] + v1[1] * v2[1]


# 创建小球
for _ in range(num_balls):
    # 随机生成小球的初始位置
    ball_pos = [random.randint(ball_radius_max, screen_width - ball_radius_max),
                random.randint(ball_radius_max, screen_height - ball_radius_max)]

    # 避免初始位置重叠
    while any(math.hypot(b['pos'][0] - ball_pos[0], b['pos'][1] - ball_pos[1]) < ball_radius_max + b['radius'] for b in
              balls):
        ball_pos = [random.randint(ball_radius_max, screen_width - ball_radius_max),
                    random.randint(ball_radius_max, screen_height - ball_radius_max)]

    # 随机生成小球的速度
    ball_speed = [random.choice(speed_range), random.choice(speed_range)]

    # 随机生成小球的半径
    ball_radius = random.randint(ball_radius_min, ball_radius_max)

    # 随机生成小球的颜色
    ball_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

    balls.append({'pos': ball_pos, 'speed': ball_speed, 'radius': ball_radius, 'color': ball_color})

# 游戏主循环
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # 更新每个小球的位置
    for i, ball_a in enumerate(balls):
        ball_a['pos'][0] += ball_a['speed'][0]
        ball_a['pos'][1] += ball_a['speed'][1]

        # 检查小球是否碰撞到窗口边界
        if ball_a['pos'][0] < 0 or ball_a['pos'][0] > screen_width - ball_a['radius']:
            ball_a['speed'][0] = -ball_a['speed'][0]
        if ball_a['pos'][1] < 0 or ball_a['pos'][1] > screen_height - ball_a['radius']:
            ball_a['speed'][1] = -ball_a['speed'][1]

        # 检查与其他小球的碰撞
        for ball_b in balls[i + 1:]:
            distance = math.hypot(ball_a['pos'][0] - ball_b['pos'][0], ball_a['pos'][1] - ball_b['pos'][1])
            min_distance = ball_a['radius'] + ball_b['radius']
            if distance < min_distance:
                calculate_new_velocities(ball_a, ball_b)

    # 填充背景颜色
    screen.fill((255, 255, 255))

    # 绘制所有小球
    for ball in balls:
        pygame.draw.circle(screen, ball['color'], ball['pos'], ball['radius'])

    # 更新屏幕显示
    pygame.display.flip()

    # 控制游戏帧率
    pygame.time.delay(10)

# 清理并退出Pygame
pygame.quit()
sys.exit()